#pragma once

#include <Eigen/Dense>
#include <vector>
#include <random>
#include <algorithm>
#include <stdexcept>

namespace  ldl_eigen
{
template <typename Scalar>
class NoCopyBatchIterator {
public:
    using MatrixType = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
    using ConstMatrixRef = Eigen::Ref<const MatrixType>;
    using MatrixMap = Eigen::Map<const MatrixType>;
    
    NoCopyBatchIterator(ConstMatrixRef features, ConstMatrixRef labels, int batch_size, bool shuffle = true)
        : m_features(features), m_labels(labels), batch_size_(batch_size), shuffle_(shuffle) {
        n_samples_ = features.rows();
        n_features_ = features.cols();
        n_labels_ = labels.cols();
        indices_.resize(n_samples_);

        // 初始化索引
        for (int i = 0; i < n_samples_; ++i) {
            indices_[i] = i;
        }
        
        if (shuffle_) {
            shuffle_indices();
        }
        
        current_pos_ = 0;
    }
    
    // 获取下一批数据（无拷贝，返回Map）
    MatrixMap next() {
        if (!has_next()) {
            throw std::out_of_range("No more batches available");
        }
        
        int start = current_pos_;
        int end = std::min(current_pos_ + batch_size_, n_samples_);
        // current_pos_ = end;
        
        // 创建指向原始数据的Map
        return MatrixMap(&m_features.coeffRef(indices_[start], 0), 
                       end - start, 
                       n_features_);
    }

    MatrixMap next_labels() {
        if (!has_next()) {
            throw std::out_of_range("No more batches available");
        }
        
        int start = current_pos_;
        int end = std::min(current_pos_ + batch_size_, n_samples_);
        // current_pos_ = end;
        
        // 创建指向原始数据的Map
        return MatrixMap(&m_labels.coeffRef(indices_[start], 0), 
                       end - start, 
                       n_labels_);
    }

    void update() {
        int end = std::min(current_pos_ + batch_size_, n_samples_);
        current_pos_ = end;
    }
    
    bool has_next() const {
        return (current_pos_ + batch_size_) < n_samples_;
    }
    
    void reset() {
        current_pos_ = 0;
        if (shuffle_) {
            shuffle_indices();
        }
    }
    
private:
    void shuffle_indices() {
        std::random_device rd;
        std::mt19937 g(rd());
        std::shuffle(indices_.begin(), indices_.end(), g);
    }
    
    ConstMatrixRef m_features{};
    ConstMatrixRef m_labels{};
    int batch_size_{};
    int n_samples_{};
    int n_features_{};
    int n_labels_{};
    bool shuffle_{};
    std::vector<int> indices_{};
    int current_pos_{};
};
}
